// Ventrue AF_COMM
#include "comm.h"
#include <net/net_namespace.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/signal.h>
#include <linux/sched/signal.h>
#include <linux/errno.h>
#include <linux/string.h>
#include <linux/stat.h>
#include <linux/dcache.h>
#include <linux/namei.h>
#include <linux/socket.h>
#include <linux/un.h>
#include <linux/fcntl.h>
#include <linux/termios.h>
#include <linux/sockios.h>
#include <linux/net.h>
#include <linux/in.h>
#include <linux/fs.h>
#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/skbuff.h>
#include <linux/netdevice.h>
#include <net/net_namespace.h>
#include <net/sock.h>
#include <net/tcp_states.h>
#include <net/af_unix.h>
#include <linux/proc_fs.h>
#include <linux/seq_file.h>
#include <net/scm.h>
#include <linux/init.h>
#include <linux/poll.h>
#include <linux/rtnetlink.h>
#include <linux/mount.h>
#include <net/checksum.h>
#include <linux/security.h>
#include <linux/freezer.h>
#include <linux/file.h>
#include <linux/hash.h>

#define PF_COMM    45

static int comm_connect(struct socket *, struct sockaddr *,
    int , int );
static int comm_init_sk(struct net * ,struct socket *, int );
static int comm_create(struct net * ,struct socket *, int , int );
static int comm_bind(struct socket *, struct sockaddr *, int );
static int comm_sendmsg(struct socket *, struct msghdr *, size_t );
static int comm_recvmsg(struct socket *, struct msghdr *, size_t , int );
static int comm_listen(struct socket *, int );
static int comm_accept(struct socket *, struct socket *, int ,  bool );
static __poll_t comm_poll(struct file *, struct socket *, poll_table *);
static int comm_release(struct socket *);
static void comm_destroy(struct sock *);
#define HASH(str) ({ \
    unsigned long hash = 5381; \
    const char *s = (str); \
    while (*s) { \
        hash = ((hash << 5) + hash) + *s++; \
    } \
    hash; \
})

static atomic_long_t comm_sock_nums = ATOMIC_LONG_INIT(0);
//自定义协议对于SOCK_DGRAM类型的操作函数集
static const struct proto_ops comm_dgram_ops = {
    .family = PF_COMM,   //协议族支持不同的套接字，但需要统一认证协议族的名称
    .owner = THIS_MODULE,
    .bind = comm_bind,
    .sendmsg = comm_sendmsg,
    .recvmsg = comm_recvmsg,
    .poll = comm_poll,
    .listen = comm_listen,
    .accept = comm_accept,
    .connect = comm_connect,
    .release = comm_release,  // 新增
};

// 实现 release 函数
static int comm_release(struct socket *sock) {
    struct sock *sk = sock->sk;
    if (sk) {
        sock_put(sk);  // 减少引用计数
        sk->sk_prot->destroy(sk);  // 调用 destroy 方法
    }
    return 0;
}

// int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen);
// 用户程序调用 connect() → 内核调用 comm_connect
static int comm_connect(struct socket *sock, struct sockaddr *uaddr,
int alen, int flags){
    struct sockaddr_comm *addr = (struct sockaddr_comm *)uaddr;
    struct comm_sock *csk = comm_sk(sock->sk);

    // 检查地址长度和名称合法性
    if (alen <= sizeof(short) || !addr->comm_family ||
        addr->comm_family != AF_COMM || addr->comm_name[0] == '\0')
        return -EINVAL;

    // 检查是否已经连接
    if (csk->name[0] != '\0') {
        return -EISCONN;
    }

    // 将目标名称存入 comm_sock 结构
    strncpy(csk->name, addr->comm_name, COMM__MAX - 1);
    csk->name[COMM__MAX - 1] = '\0';
    if (csk->name[COMM__MAX - 1] != '\0') {
        return -EINVAL;
    }

    // 设置套接字状态为已连接
    sock->state = SS_CONNECTED;

    return 0;

}

//指向这个结构的指针 在socket_register()时注册到全局表
static const struct net_proto_family comm_family_ops = {
    .family = PF_COMM,
    .create = comm_create,  //用户态socket()创建套接字时起作用的内核函数
    .owner = THIS_MODULE,
};

//proto 描述一个网络协议的相关属性和操作
static struct proto comm_proto = {
	.name			= "COMM",
	.owner			= THIS_MODULE,
	.obj_size		= sizeof(struct comm_sock),//sk_alloc分配内存用到
    .destroy     = comm_destroy,  // 新增销毁方法
};

static void comm_destroy(struct sock *sk) {
    struct comm_sock *csk = comm_sk(sk);

    // 从哈希表中移除节点
    write_lock(&comm_table_lock);
    if (!hlist_unhashed(&csk->node)) {
        hlist_del_init(&csk->node);
    }
    write_unlock(&comm_table_lock);

    // 释放接收队列中的 skb
    skb_queue_purge(&csk->recv_queue);
}

//create过程主要填充struct socket和struct sock的内容
static int comm_init_sk(struct net *net ,struct socket *sock, int kern){
    printk("初始化进入到二阶段函数\n");
    struct sock *sk = NULL;
    struct comm_sock *cm;
    atomic_long_inc(&comm_sock_nums);//增加套接字数, 预防资源耗尽
	if (atomic_long_read(&comm_sock_nums) > 2 * get_max_files())
		goto OUT;
    //分配sock结构体,看似是初始化sock,其实是初始化了comm_sock的大小
    sk = sk_alloc(net, PF_COMM, GFP_KERNEL, &comm_proto, kern);
    printk("sk初始化的地址%px\n",(void *)sk);
	if (!sk){
        printk(KERN_ERR "comm: sk_alloc failed\n");
	    return -ENOMEM;
    }
    //sock和sk关联
    sock_init_data(sock,sk);
    printk("sock->sk的地址为%px\n",(void*)sock->sk);
    printk("sk->sndbuf=%d  sk->rcvbuf=%d\n",sk->sk_sndbuf ,sk->sk_rcvbuf);
    sock_hold(sk);  // 增加引用计数
    sk->sk_allocation	= GFP_KERNEL_ACCOUNT;
	sk->sk_max_ack_backlog	= net->unx.sysctl_max_dgram_qlen;
	//sk->sk_destruct		= comm_destruct;//套接字销毁时的回调

    //初始化comm_sock结构体的各个字段
    // 将结构体进行转换 cm扩充sk
    cm = comm_sk(sk);
    cm->name[0] = '\0';
    int i=0;
    for (i = 0; i < 256; i++) {
        INIT_HLIST_HEAD(&comm_table[i]);
    }
    // INIT_HLIST_NODE(&cm->node);
    spin_lock_init(&cm->lock);
    skb_queue_head_init(&cm->recv_queue);
    printk("初始化完成\n");
    return sk;

OUT:
    atomic_long_dec(&comm_sock_nums);
    printk("初始化进入到OUT\n");
    return -ENAVAIL;
}

//int socket(int domain, int type, int protocol)
//用户调用 socket(AF_COMM, SOCK_DGRAM, 0) → 内核调用 comm_create
//仿写UNIX域 kern为1:内核socket 为0:用户态socket
static int comm_create(struct net *net ,struct socket *sock, int protocol, int kern){
    //protocol检查
    if(protocol && protocol != PF_COMM){
        return -EPROTONOSUPPORT;
    }

    //type检查
    switch (sock->type){
    case SOCK_DGRAM:
        sock->ops = &comm_dgram_ops;
        break;
    case SOCK_STREAM:
    case SOCK_SEQPACKET:
    case SOCK_RAW:
        //暂不支持
	    break;
    default:
        return -ESOCKTNOSUPPORT;
    }

    //设定socket状态未连接
    sock->state = SS_UNCONNECTED;
    int ret = comm_init_sk(net ,sock, kern)? 0 : -ENOMEM;
    printk("ret=%d\n", ret);
    return ret;
}


//int bind(int sockfd, const struct sockaddr *addr, socklen_t addrlen)
// 用户程序调用 bind() → 内核调用 comm_bind
static int comm_bind(struct socket *sock, struct sockaddr *uaddr, int addr_len){
    printk("准备执行comm_bind\n");
    printk("检查入参地址%px   %px\n",(void *)sock ,(void *)uaddr);
    struct sockaddr_comm *addr = (struct sockaddr_comm *)uaddr;

    struct comm_sock *csk = comm_sk(sock->sk);
    printk("sock->sk=%px  \n",(void *)sock->sk );
    // 检查地址长度和名称合法性
    if (addr_len <= sizeof(short) || !addr->comm_family ||
        addr->comm_family != AF_COMM || addr->comm_name[0] == '\0')
        return -EINVAL;

    // 将名称存入 comm_sock 结构
    if (csk->name[0] != '\0') {
        printk("csk->name不为空 返回\n");
        return -EADDRINUSE;
    }
    strncpy(csk->name, addr->comm_name, COMM__MAX -1);
    csk->name[COMM__MAX -1 ] = '\0';
    if(csk->name[COMM__MAX - 1] != '\0'){
        printk("csk->name未正常终止 返回\n");
        return -EINVAL;
    }

    // 将socket节点 添加到全局哈希表
    // 不用hash_add() 其不允许同名节点
    printk("即将获取哈希表锁\n");
    write_lock(&comm_table_lock);
    //hash_add(comm_table, &csk->node, hash_string(csk->name));
    unsigned int hash = HASH(csk->name) % 256;
    printk("插入哈希表: name=%s, hash=%lu\n", csk->name, HASH(csk->name)%256);
    struct hlist_head *head = &comm_table[hash];
    hlist_add_head(&csk->node, head);

    printk("插入后验证哈希桶 %d:\n", hash);
    int count = 0;
    struct comm_sock *tmp;
    hlist_for_each_entry(tmp, head, node) {
        count++;
        printk("发现节点: name=%s (计数: %d)\n", tmp->name, count);
    }
    printk("哈希桶 %d 中共有 %d 个节点\n", hash, count);

    write_unlock(&comm_table_lock);
    printk("comm_bind正确返回\n");
    return 0;
}

static int comm_sendmsg(struct socket *sock, struct msghdr *msg, size_t len){
    struct comm_sock *csk = comm_sk(sock->sk);
    struct sk_buff *skb;
    struct comm_sock *target;
    int err = 0;
    printk("准备发送信息comm_sendmsg\n");
    // 克隆 skb 数据报
    skb = sock_alloc_send_skb(sock->sk, len, 0, &err);
    if(!skb){
        printk("skb为空 返回\n");
        return err;
    }
    err = memcpy_from_msg(skb_put(skb, len), msg, len);
    printk("成功拷贝msg\n");
    if (err) {
        printk("拷贝失败 返回\n");
        kfree_skb(skb);
        return err;
    }
    
    unsigned int hash = HASH(csk->name) % 256;  // 必须与bind时相同的取模逻辑
    struct hlist_head *head = &comm_table[hash]; // 定位哈希桶
    printk("发送方名称: %s, 哈希桶: %u\n", csk->name, hash);
    // 遍历哈希表 广播到所有同名socket
    read_lock(&comm_table_lock);
    hlist_for_each_entry(target, head, node) {  // 直接遍历该桶的链表
        printk("遍历到节点: name=%s\n", target->name);
        if (strcmp(target->name, csk->name) == 0) { 
            // 4. 克隆skb并广播给同名socket
            struct sk_buff *skb_cln = skb_clone(skb, GFP_ATOMIC);
            if (skb_cln) {
                spin_lock(&target->lock);
                skb_queue_tail(&target->recv_queue, skb_cln);
                spin_unlock(&target->lock);
                printk("成功发送到: %s\n", target->name);
            } else {
                printk("克隆skb失败\n");
            }
        }
    }
    read_unlock(&comm_table_lock);


    // 释放原始 skb
    consume_skb(skb);
    printk("sendmsg正确返回\n");
    return len;
}

static int comm_recvmsg(struct socket *sock, struct msghdr *msg, size_t len, int flags){
    struct comm_sock *csk = comm_sk(sock->sk);
    struct sk_buff *skb;
    int err = 0;
    printk("准备接收信息comm_recvmsg\n");
    // 从接收队列中取出skb
    spin_lock(&csk->lock);
    skb = skb_dequeue(&csk->recv_queue);
    spin_unlock(&csk->lock);

    if (!skb) {
        // 队列为空，返回错误或根据标志位处理
        if (flags & MSG_DONTWAIT) {
            printk("队列为空，返回错误\n");
            return -EAGAIN;
        }
        // TODO:这里可以添加阻塞等待逻辑
        printk("skb为空 返回\n");
        return -EAGAIN;
    }

    if (skb->len < len) {
        len = skb->len;
    }

    // 将数据从skb复制到用户空间
    err = memcpy_to_msg(msg, skb->data, len);
    if (err) {
        // 复制失败 释放skb
        kfree_skb(skb);
        printk("复制skb失败 返回\n");
        return err;
    }

    // 释放skb
    consume_skb(skb);
    printk("recvmsg正确返回\n");
    return len;
}


static int comm_listen(struct socket *sock, int backlog){
    // 数据报套接字不支持 listen 操作
    return -EOPNOTSUPP;
}

static int comm_accept(struct socket *sock, struct socket *newsock, int flags,
    bool kern){
    // 数据报套接字不支持 accept 操作
    return -EOPNOTSUPP;
}
//当用户态传入多个FD时 内核遍历所有的Fd然后分别执行poll函数  一次性完成所有FD的检验
static __poll_t comm_poll(struct file *file, struct socket *sock, poll_table *wait) {
    struct sock *sk = sock->sk;
    struct comm_sock *csk = comm_sk(sk);
    __poll_t mask = 0;

    sock_poll_wait(file, sock, wait);//当前进程加入到socket套接字的等待队列中

    // 异常事件处理
    if (sk->sk_err || !skb_queue_empty_lockless(&sk->sk_error_queue))
        mask |= EPOLLERR;

    // 可读事件
    if (!skb_queue_empty_lockless(&csk->recv_queue))
        mask |= EPOLLIN | EPOLLRDNORM;

    // 关闭事件
    if (sk->sk_shutdown & RCV_SHUTDOWN)
        mask |= EPOLLRDHUP | EPOLLIN | EPOLLRDNORM;
    if (sk->sk_shutdown == SHUTDOWN_MASK)
        mask |= EPOLLHUP;

    return mask;
}

static int __init comm_init(void){
    //注册完之后,将comm_family_ops赋值给了全局指针net_families
    // sock_register(&comm_family_ops);
    // return 0;
    return sock_register(&comm_family_ops);
}

static void __exit comm_exit(void) {
    sock_unregister(PF_COMM);
}

//fs_initcall(comm_init);//内核内置功能时采用
module_init(comm_init);//可动态加载模块LKM的初始化(即通过insmod/modprobe加载时)
module_exit(comm_exit);

MODULE_LICENSE("GPL");//使用module_init时需声明该许可 GPL才能与内核互动
